//
// Created by mio on 25-4-23.
//

#include <QDebug>

#include "RtmpConnection.h"
#include "rtmp.h"
#include "RtmpPublisher.h"

RtmpConnection::RtmpConnection(std::shared_ptr<RtmpPublisher> publisher, TaskScheduler *scheduler, int sockfd)
    :RtmpConnection(scheduler,sockfd,publisher.get())
{
    handshake_.reset(new RtmpHandshake(RtmpHandshake::HANDSHAKE_S0S1S2));
    rtmp_publisher_ = publisher;
}

RtmpConnection::RtmpConnection(TaskScheduler* scheduler, int sockfd, Rtmp* rtmp)
    : TcpConnection(scheduler, sockfd),
    rtmp_chunk_(new RtmpChunk()),
    state_(HANDSHAKE)
{
    amf_decoder_.reset(new AmfDecoder());
    amf_encoder_.reset(new AmfEncoder());

    max_chunk_size_ = rtmp->GetChunkSize();
    stream_path_ = rtmp->GetStreamPath();
    stream_name_ = rtmp->GetStreamName();
    app_ = rtmp->GetApp();

    this->SetReadCallback([this](std::shared_ptr<TcpConnection> conn, BufferReader& buffer)
    {
        return this->OnRead(buffer);
    });

    this->SetCloseCallback([this](std::shared_ptr<TcpConnection> conn)
    {
        this->OnClose();
    });
}

RtmpConnection::~RtmpConnection()
{
}

bool RtmpConnection::OnRead(BufferReader& buffer)
{
    bool ret = true;
    if (handshake_->IsCompleted()) // 是否握手完成, 完成之后才能发送消息
    {
        ret = HandleChunk(buffer);
    }
    else
    {
        std::shared_ptr<char> res(new char[4096], std::default_delete<char[]>());
        int res_size = handshake_->Parse(buffer, res.get(), 4096);
        if (res_size < 0) ret = false;
        if (res_size > 0) this->Send(res.get(), res_size);
        if (handshake_->IsCompleted())
        {
            // 如果握手刚完成，并且缓冲区中还有剩余数据，尝试处理这些数据
            if (buffer.ReadableBytes() > 0)
                ret = HandleChunk(buffer);

            // TODO: 发送连接请求
            SetChunkSize();
            Connect();
        }
    }
    return ret;
}

void RtmpConnection::OnClose()
{
    this->DeleteStream();
}

bool RtmpConnection::HandleChunk(BufferReader& buffer)
{
    int ret = -1;
    do
    {
        RtmpMessage rtmp_msg;
        ret = rtmp_chunk_->Parse(buffer, rtmp_msg);
        if (ret >= 0) // 解析成功, 需要处理message
        {
            if (rtmp_msg.IsCompleted())
            {
                if (!HandleMessage(rtmp_msg)) return false;
            }
            if (ret == 0) break; // 缓冲区没有数据
        }
        else return false;
    } while (buffer.ReadableBytes() > 0);
    return true;
}

bool RtmpConnection::HandleMessage(RtmpMessage& rtmp_msg)
{
    bool ret = true;
    switch (rtmp_msg.type_id)
    {
    case RTMP_INVOKE:    // 视频消息 (Type ID 9)
        ret = HandleInvoke(rtmp_msg);
        break;
    case RTMP_SET_CHUNK_SIZE:    // 音频消息 (Type ID 8)
        rtmp_chunk_->SetInChunkSize(ReadUint32BE(rtmp_msg.payload.get()));
        break;
    default:
        break;
    }
    return ret;
}

bool RtmpConnection::Handshake()
{
    uint32_t size = 1 + 1536;
    std::shared_ptr<char> req(new char[size], std::default_delete<char[]>());
    handshake_->BuildC0C1(req.get(), size);
    this->Send(req.get(), size);
    return true;
}

bool RtmpConnection::Connect()
{
    // TODO: 把number改为0
    AmfObjects objects;
    amf_encoder_->reset();
    amf_encoder_->encodeString("connect", 7);
    amf_encoder_->encodeNumber((double)++number_);
    objects["app"] = AmfObject(app_);
    objects["type"] = AmfObject(std::string("nonprivate"));
    amf_encoder_->encodeObjects(objects);

    // 更新状态
    state_ = START_CONNECT;
    SendInvokeMsg(RTMP_CHUNK_INVOKE_ID, amf_encoder_->data(), amf_encoder_->size());
    qDebug() << "Connected";
    return true;
}

bool RtmpConnection::CretaeStream()
{
    amf_encoder_->reset();
    amf_encoder_->encodeString("createStream", 12);
    amf_encoder_->encodeNumber((double)++number_);
    AmfObjects objects;
    amf_encoder_->encodeObjects(objects);

    // 更新状态
    state_ = START_CREATE_STREAM;
    SendInvokeMsg(RTMP_CHUNK_INVOKE_ID, amf_encoder_->data(), amf_encoder_->size());
    qDebug() << "Connected";
    return true;
}

bool RtmpConnection::Publish()
{
    amf_encoder_->reset();
    amf_encoder_->encodeString("publish", 7);
    amf_encoder_->encodeNumber((double)++number_);
    AmfObjects objects;
    amf_encoder_->encodeObjects(objects);
    amf_encoder_->encodeString(stream_name_.c_str(), (int)stream_name_.size());
    // amf_encoder_->encodeString("live", 4);

    //更新状态
    state_ = START_PUBLISH;
    SendInvokeMsg(RTMP_CHUNK_INVOKE_ID,amf_encoder_->data(),amf_encoder_->size());
    return true;
}

bool RtmpConnection::DeleteStream()
{
    //编码amf
    AmfObjects objects;
    amf_encoder_->reset();

    amf_encoder_->encodeString("DeleteStream",12);
    amf_encoder_->encodeNumber((double)++number_);
    amf_encoder_->encodeObjects(objects);
    amf_encoder_->encodeNumber(stream_id_);

    //更新状态
    state_ = START_DELETE_STREAM;
    SendInvokeMsg(RTMP_CHUNK_INVOKE_ID,amf_encoder_->data(),amf_encoder_->size());
    return true;
}


bool RtmpConnection::HandleInvoke(RtmpMessage& rtmp_msg)
{
    bool ret = true;
    amf_decoder_->reset();
    int byte_used = amf_decoder_->decode((const char*)rtmp_msg.payload.get(), rtmp_msg.length, 1);
    if (byte_used < 0) { return false; }

    std::string method = amf_decoder_->getString();
    byte_used = amf_decoder_->decode(rtmp_msg.payload.get() + byte_used, rtmp_msg.length - byte_used);

    if (method == "_result") { ret = HandleResult(rtmp_msg); }
    else if (method == "onStatus"){ ret = HandleOnStatus(rtmp_msg); }
    return ret;
}

bool RtmpConnection::HandleResult(RtmpMessage &rtmp_msg)
{
    bool ret = false;
    if (state_ == START_CONNECT)
    {
        if (amf_decoder_->hasObject("code"))
        {
            AmfObject amf_object = amf_decoder_->getObject("code");
            if (amf_object.amf_string == "NetConnection.Connect.Success")
            {
                qDebug() << "Start Connect";
                CretaeStream();
                ret = true;
            }
        }
    }
    else if (state_ == START_CREATE_STREAM) {
        if (amf_decoder_->getNumber() > 0) {
            stream_id_ = (int)amf_decoder_->getNumber();
            qDebug() << "START_CREATE_STREAM";
            Publish();
            ret = true;
        }
    }
    return ret;
}

bool RtmpConnection::HandleOnStatus(RtmpMessage &rtmp_msg)
{
    bool ret = true;
    if (state_ == START_PUBLISH) {
        if (amf_decoder_->hasObject("code")) {
            AmfObject amf_object = amf_decoder_->getObject("code");
            std::string code = amf_object.amf_string;
            if (code == "NetStream.Publish.Start") is_publishing_ = true;
            else ret = false;
        }
    }
    else if (state_ == START_DELETE_STREAM) {
        if (amf_decoder_->hasObject("code")) {
            AmfObject amf_object = amf_decoder_->getObject("code");
            std::string code = amf_object.amf_string;
            if (code != "NetStream.Unpublish.Success") ret = false; // 取消推流失败
        }
    }
    return ret;
}


void RtmpConnection::SetChunkSize()
{
    rtmp_chunk_->SetOutChunkSize(max_chunk_size_);
    std::shared_ptr<char> data(new char[4], std::default_delete<char[]>());
    WriteUint32BE(data.get(), max_chunk_size_);
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_SET_CHUNK_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 4;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

bool RtmpConnection::SendInvokeMsg(uint32_t csid, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (this->IsClosed()) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_INVOKE;
    rtmp_msg.timestamp = 0;
    rtmp_msg.length = payload_size;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    SendRtmpChunks(csid, rtmp_msg);
    return true;
}

bool RtmpConnection::SendNotifyMsg(uint32_t csid, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (this->IsClosed()) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_NOTIFY;
    rtmp_msg.timestamp = 0;
    rtmp_msg.length = payload_size;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    SendRtmpChunks(csid, rtmp_msg);
    return true;
}

bool RtmpConnection::IsKeyFrame(std::shared_ptr<char> data, uint32_t size)
{
    // 从第一个字节获取帧类型和编码 ID
    uint8_t frame_type = (data.get()[0] >> 4) & 0x0f;
    uint8_t code_id = data.get()[0] & 0x0f;
    return (frame_type == 1 && code_id == RTMP_CODEC_ID_H264);
}

bool RtmpConnection::SendVideoData(uint64_t timestamp, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (payload_size <= 0) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_VIDEO;
    rtmp_msg.timestamp = timestamp;
    rtmp_msg.payload = payload;
    rtmp_msg.length = payload_size;

    SendRtmpChunks(RTMP_CHUNK_VIDEO_ID, rtmp_msg);
    return true;
}

bool RtmpConnection::SendAudioData(uint64_t timestamp, std::shared_ptr<char> payload, uint32_t payload_size)
{
    if (payload_size <= 0) return false;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_AUDIO;
    rtmp_msg.timestamp = timestamp;
    rtmp_msg.payload = payload;
    rtmp_msg.length = payload_size;

    SendRtmpChunks(RTMP_CHUNK_AUDIO_ID, rtmp_msg);
    return true;
}

void RtmpConnection::SendRtmpChunks(uint32_t csid, RtmpMessage& rtmp_msg)
{
    // 粗略估算需要的缓冲区大小：消息长度 + 每个块可能的头开销 + 额外余量
    uint32_t capacity = rtmp_msg.length + rtmp_msg.length / max_chunk_size_ * 5 + 1024;
    std::shared_ptr<char> buffer(new char[capacity], std::default_delete<char[]>());
    int size = rtmp_chunk_->CreateChunk(csid, rtmp_msg, buffer.get(), capacity);
    if (size > 0)
    {
        this->Send(buffer.get(), size);
    }
}
